from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    utility,
)
from sentence_transformers import SentenceTransformer
from typing import List, Union

# 关于 paraphrase-multilingual-MiniLM-L12-v2 这种embedding模型
# 程序运行会自动检测模型文件在本机是否存在，如果不存在会自动从huggingface上下载。
# 默认存放在用户目录下的 .cache 文件夹中，例如：C:\Users\username\.cache\huggingface\hub
# 如果不想保存在C盘，可以考虑使用软连接指向其他盘符。注意执行此命令之前，目标文件夹（后者）可以存在也可以不存在，来源文件夹（前者）必须不存在。(myfox是我的用户名，你要改成你自己的)
# 例如我这里的：mklink /J "C:\Users\myfox\.cache\huggingface\" "E:\ProgramFiles\AIapplication\huggingface_cache\"


class VectorDatabase:
    """向量数据库操作封装类"""

    def __init__(
        self,
        host="localhost",
        port="19530",
        model_name="paraphrase-multilingual-MiniLM-L12-v2",
    ):
        """初始化向量数据库连接和模型

        Args:
            host: Milvus服务地址
            port: Milvus服务端口
            model_name: 使用的embedding模型名称
        """
        self.host = host
        self.port = port
        self.model = SentenceTransformer(model_name)
        self.collection = None

    def connect(self):
        """连接Milvus数据库"""
        try:
            connections.connect(
                host=self.host, port=self.port, timeout=10, secure=False
            )
            print(
                f"Milvus连接成功！服务地址:{connections.get_connection_addr('default').get('address')}"
            )
            return True
        except Exception as e:
            print(f"连接失败: {e}")
            return False

    def create_collection(
        self, collection_name, fields: List[FieldSchema], description=""
    ):
        """创建及选择集合（如果不存在则创建，如果存在则直接选择）

        Args:
            collection_name: 集合名称
            fields: 字段结构
            description: 集合描述
        Returns:
            bool: 是否创建了新集合
        """
        # 检查集合是否已存在
        if collection_name in utility.list_collections():
            print(f"集合 {collection_name} 已存在，无需创建")
            self.collection = Collection(collection_name)
            return False

        schema = CollectionSchema(fields, description=description)
        self.collection = Collection(collection_name, schema)
        print(f"新建集合了: {collection_name}")
        print(f"现在的集合列表: {utility.list_collections()}")
        return True

    def drop_collection(self, collection_name):
        """删除集合，可以处理本来就不存在的情况

        Args:
            collection_name: 要删除的集合名称
        """
        if collection_name in utility.list_collections():
            utility.drop_collection(collection_name)
            print(f"已删除集合: {collection_name}")
        else:
            print(f"集合 {collection_name} 现在不存在，无需删除")

    def insert_data(self, vector_mapping=None, **kwargs):
        """插入向量数据(支持字段映射)

        Args:
            vector_mapping: 向量化字段映射字典，格式为:
                          {"目标向量字段": "源文本字段"}
                          例如: {"vector": "description"}
            **kwargs: 字段数据，键为字段名，值为字段值列表
                    示例: title=["商品1", "商品2"],
                         description=["描述1", "描述2"]
        Returns:
            插入数据的primary_key列表
        """
        if not self.collection:
            print("请先创建或选择集合")
            return None

        if not vector_mapping or not kwargs:
            print("错误: 必须提供vector_mapping和字段数据")
            return None

        # 准备插入数据
        data = []
        # 处理源字段
        for field in self.collection.schema.fields:
            if field.name in kwargs:
                data.append(kwargs[field.name])

        # 处理向量化字段
        for vector_field, text_field in vector_mapping.items():
            print(f"处理向量化字段: {text_field} -> {vector_field}")
            if text_field in kwargs:
                vectors = self.model.encode(kwargs[text_field]).tolist()
                data.append(vectors)
            else:
                print(f"错误: 缺少向量化字段 {text_field} 的文本数据")
                return None

        # print("准备插入的数据:", data)
        # Milvus的插入数据要求字段顺序与集合定义一致
        # Milvus的主键用于标识而没有约束，所以反复插入相同的主键不会报错
        # 因此，如果我们需要按照主键控制数据的唯一性，需要在插入时先查询再处理，但会增加查询的开销。
        insert_result = self.collection.insert(data)
        print("插入数据主键:", insert_result.primary_keys)
        return insert_result.primary_keys

    def create_index(
        self,
        field_name: Union[List[str], str],
        index_type="IVF_FLAT",
        metric_type="L2",
        nlist=128,
    ):
        """创建索引(支持多字段)

        Args:
            field_name: 需要创建索引的字段名列表，如["vector", "title_vector"]，也可以传入单个字段字符串或列表。
            index_type: 索引类型
            metric_type: 相似度度量类型
            nlist: 索引参数
        """
        if not self.collection:
            print("请先创建或选择集合")
            return

        # 确保field_name是列表形式
        if isinstance(field_name, str):
            field_name = [field_name]

        index_params = {
            "index_type": index_type,
            "metric_type": metric_type,
            "params": {"nlist": nlist},
        }

        # 为每个字段创建索引
        for field in field_name:
            self.collection.create_index(field_name=field, index_params=index_params)
            print(f"字段 {field} 索引构建完成")

        # 加载集合以应用索引
        self.collection.load()

    def search(
        self,
        query_text,
        vector_field,
        output_fields,
        limit=3,
        nprobe=10,
        metric_type="L2",
    ):
        """执行相似度搜索

        Args:
            query_text: 查询文本
            vector_field: 向量字段名，字符串，如"vector"
            output_fields: 需要返回的字段列表，如["title", "description"]
            limit: 返回结果数量
            nprobe: 搜索参数
        Returns:
            搜索结果列表
        """
        if not self.collection:
            print("请先创建或选择集合")
            return None

        # 需要刷新集合，否则获取num_entities可能不准确，是非实时的
        self.collection.flush()
        print(f"集合中现有数据量: {self.collection.num_entities}")

        query_vector = self.model.encode(query_text)
        search_params = {"metric_type": metric_type, "params": {"nprobe": nprobe}}

        results = self.collection.search(
            data=[query_vector.tolist()],
            anns_field=vector_field,
            param=search_params,
            limit=limit,
            output_fields=output_fields,
        )

        # 格式化搜索的结果
        search_results = []
        for rank, hit in enumerate(results[0]):
            search_result = {
                "rank": rank + 1,
                "distance": hit.distance,
                "similarity": 1 - hit.distance,  # 计算相似度，值越大相似度越高
            }
            for output_field in output_fields:
                search_result[output_field] = hit.entity.get(output_field)

            search_results.append(search_result)

        return search_results


if __name__ == "__main__":
    # 示例用法
    db = VectorDatabase()
    if db.connect():
        # 删除集合
        db.drop_collection("products")
        # 建议先删除再重新构建，因为Milvus的主键用于标识而没有约束，所以反复插入相同的主键不会报错，会造成数据重复。

        # 创建集合（如果存在则直接选择）
        fields = [
            FieldSchema(
                name="title",
                dtype=DataType.VARCHAR,
                is_primary=True,
                max_length=100,
            ),
            FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=500),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=384),
        ]
        db.create_collection("products", fields=fields, description="商品特征库")

        # 插入数据
        titles = [
            "智能手机X20",
            "无线降噪耳机Max",
            "4K超高清智能电视",
            "迷你便携式投影仪",
            "智能健康手表S7",
            "机械背光键盘G90",
            "光学游戏鼠标M5",
            "2K屏平板电脑P8",
            "彩色家用打印机P30",
            "虚拟环绕声电竞耳机H9",
        ]
        descriptions = [
            "搭载顶级处理器，具备超强性能和流畅体验的智能手机，支持5G网络和快速充电功能",
            "采用先进降噪技术，提供沉浸式音乐体验和清晰通话质量的无线蓝牙耳机",
            "拥有4K超高清分辨率和智能交互系统，带来震撼视觉效果的智能电视",
            "体积小巧轻便，便于携带和随时随地投影使用的迷你投影仪",
            "集成多种健康监测功能，如心率、睡眠监测等的多功能智能手表",
            "配备背光设计，按键手感舒适且反馈灵敏的机械键盘，适合游戏和办公使用",
            "采用高精度光学传感器，实现精准操控和快速响应的游戏鼠标",
            "配备2K高清屏幕，轻薄便携，适合办公、娱乐和学习的平板电脑",
            "支持彩色打印，具备高效打印速度和出色打印质量的家用打印机",
            "采用虚拟环绕声技术，提供沉浸式游戏音效体验的电竞耳机",
        ]
        db.insert_data(
            vector_mapping={"vector": "description"},
            title=titles,
            description=descriptions,
        )

        # 创建索引
        db.create_index(field_name="vector")

        # 执行搜索
        search_results = db.search(
            "推荐一个降噪比较好的耳机",
            vector_field="vector",
            output_fields=["title", "description"],
        )
        for search_result in search_results:
            print(
                f"[搜索结果] 第{search_result['rank']}名: {search_result['title']}，向量距离：{search_result['distance']}，商品描述：{search_result['description']}，相似度: {search_result['similarity']:.2f}"
            )

        # 删除集合(可选)
        # db.drop_collection("products")
